﻿#pragma once
#include <utility>
#include <assert.h>
#include <iostream>

using namespace std;

template<class K, class V>
class treeNode
{
	typedef treeNode<K, V> Node;
	template<class K, class V>
	friend class AVLTree;
public:
	treeNode(const pair<K, V>& kv)
		:_kv(kv)
	{}
	treeNode()
	{}
private:
	Node* _left = nullptr;
	Node* _right = nullptr;
	Node* _parent = nullptr;
	pair<K, V> _kv;
	int _bf = 0;
};

template<class K, class V>
class AVLTree
{
	typedef treeNode<K, V> Node;
public:
	bool insert(const pair<K, V>& kv)
	{
		if (root == nullptr)
		{
			root = (Node*)new Node(kv);
			return true;
		}
		Node* cut = root;
		Node* rem = nullptr;
		while (cut != nullptr)
		{
			if (cut->_kv.first > kv.first)
			{
				rem = cut;
				cut = cut->_left;
			}
			else if (cut->_kv.first < kv.first)
			{
				rem = cut;
				cut = cut->_right;
			}
			else
			{
				return false;
			}
		}
		cut = (Node*)new Node(kv);
		if (rem->_kv.first > cut->_kv.first)
		{
			rem->_left = cut;
			rem->_bf--;
		}
		else
		{
			rem->_right = cut;
			rem->_bf++;
		}
		cut->_parent = rem;
		while (rem != nullptr)
		{
			if (rem->_bf == 0)
				return true;
			else if (rem->_bf == -1 || rem->_bf == 1)
			{
				cut = rem;
				rem = rem->_parent;
				if (rem == nullptr)
					return true;
				if (rem->_left == cut)
					rem->_bf--;
				else
					rem->_bf++;
			}
			else if (rem->_bf == 2 || rem->_bf == -2)
			{
				if (rem->_bf == -2 && cut->_bf == -1)
					rotateR(rem);
				else if (rem->_bf == 2 && cut->_bf == 1)
					rotateL(rem);
				else if (rem->_bf == -2 && cut->_bf == 1)
					rotateLR(rem);
				else if (rem->_bf == 2 && cut->_bf == -1)
					rotateRL(rem);
				else
					assert(false);
				break;
			}
			else
			{
				assert(false);
			}
		}

		return true;
	}

	void inoder()
	{
		_inoder(root);
	}

	int getHeight(Node* root)
	{
		if (root == nullptr)
		{
			return 0;
		}
		int left = getHeight(root->_left);
		int right = getHeight(root->_right);
		return 1 + (left > right ? left : right);
	}

	bool checkAVL()
	{
		return _checkAVL(this->root);
	}

	

private:

	bool _checkAVL(Node* root)
	{
		if (root == nullptr)
			return true;
		int left = getHeight(root->_left);
		int right = getHeight(root->_right);
		int diff = right - left;
		if (abs(diff) >= 2)
			return false;
		if (diff != root->_bf)
			return false;
		return _checkAVL(root->_left) && _checkAVL(root->_right);
	}

	void _inoder(Node* root)
	{
		if (root == nullptr)
			return;
		_inoder(root->_left);
		cout << root->_kv.first << ":" << root->_kv.second << endl;
		_inoder(root->_right);
		return;
	}

	//ҵ
	void rotateR(Node*& par)
	{
		Node* chi = par->_left;
		Node* chiR = chi->_right;
		par->_left = chiR;
		if (chiR)
			chiR->_parent = par;
		chi->_right = par;
		chi->_parent = par->_parent;
		par->_parent = chi;
		if (chi->_parent)
		{
			if (chi->_parent->_left == par)
				chi->_parent->_left = chi;
			else
				chi->_parent->_right = chi;
		}
		if (root == par)
			root = chi;
		par->_bf = 0;
		chi->_bf = 0;
	}
	//
	void rotateL(Node*& par)
	{
		Node* chi = par->_right;
		Node* chiL = chi->_left;
		par->_right = chiL;
		if (chiL != nullptr)
			chiL->_parent = par;
		chi->_left = par;
		chi->_parent = par->_parent;
		par->_parent = chi;
		if (chi->_parent)
		{
			if(chi->_parent->_left == par)
				chi->_parent->_left = chi;
			else
				chi->_parent->_right = chi;
		}
		if (root == par)
			root = chi;
		par->_bf = 0;
		chi->_bf = 0;
	}
	//˫
	void rotateLR(Node*& par)
	{
		Node* chi = par->_left;
		Node* chiR = chi->_right;
		int bf = chiR->_bf;
		rotateL(par->_left);
		rotateR(par);
		if (bf == -1)
		{
			par->_bf = 1;
			chi->_bf = 0;
			chiR->_bf = 0;
		}
		else if (bf == 1)
		{
			chi->_bf = -1;
			par->_bf = 0;
			chiR->_bf = 0;
		}
		else if (bf == 0)
		{
			chi->_bf = 0;
			par->_bf = 0;
			chiR->_bf = 0;
		}
		else
		{
			assert(false);
		}
	}
	void rotateRL(Node*& par)
	{
		Node* chi = par->_right;
		Node* chiL = chi->_left;
		int bf = chiL->_bf;
		rotateR(par->_right);
		rotateL(par);
		if (bf == -1)
		{
			par->_bf = -1;
			chi->_bf = 0;
			chiL->_bf = 0;
		}
		else if (bf == 1)
		{
			chi->_bf = 1;
			par->_bf = 0;
			chiL->_bf = 0;
		}
		else if (bf == 0)
		{
			chi->_bf = 0;
			par->_bf = 0;
			chiL->_bf = 0;
		}
		else
		{
			assert(false);
		}

	}

	


	Node* root = nullptr;
};